package guhong.play.commandsystem.job.tool.ai.service;

import cn.hutool.core.util.StrUtil;
import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationParam;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.common.ResultCallback;
import com.alibaba.dashscope.common.Role;
import com.alibaba.dashscope.exception.ApiException;
import com.alibaba.dashscope.exception.InputRequiredException;
import com.alibaba.dashscope.exception.NoApiKeyException;
import com.alibaba.fastjson.JSONObject;
import guhong.play.commandsystem.CommandManager;
import guhong.play.commandsystem.job.tool.ai.AiCommandManager;
import guhong.play.commandsystem.job.tool.ai.AiConfig;
import guhong.play.commandsystem.util.print.PrintUtil;
import lombok.Data;

import java.util.Arrays;
import java.util.List;
import java.util.concurrent.Semaphore;

/**
 * 通义千问
 *
 * @author : 李双凯
 * @date : 2019-11-20 22:32
 **/
@Data
public class QwenAiService implements AiService {


    @Override
    public void getAnswer(String messageStr) {
        try {

            Generation gen = new Generation();

            AiCommandManager.appendMessage(messageStr);

            PrintUtil.println();
            List<Message> genMessageList = getMessageList();
            streamCallWithCallback(gen, genMessageList);
        } catch (ApiException | NoApiKeyException | InputRequiredException | InterruptedException e) {
            PrintUtil.printLnError("An exception occurred: " + e.getMessage());
        }
    }

    // 写一个获取apiKey的方法
    public static String getApiKey() {
        JSONObject thirdParty = CommandManager.getSystemConfig().getThirdParty();
        JSONObject aliyunConfig = thirdParty.getJSONObject("aliyun");
        JSONObject aiConfig = aliyunConfig.getJSONObject("ai");
        String apiKey = aiConfig.getString("apiKey");
        if (StrUtil.isBlank(apiKey)) {
            throw new RuntimeException("请先配置阿里云的apiKey");
        }
        return apiKey;
    }



    public static void streamCallWithCallback(Generation gen, List<Message> userMsgList)
            throws NoApiKeyException, ApiException, InputRequiredException, InterruptedException {
        GenerationParam param = buildGenerationParam(userMsgList);
        Semaphore semaphore = new Semaphore(0);

        gen.streamCall(param, new ResultCallback<GenerationResult>() {
            @Override
            public void onEvent(GenerationResult message) {
                // 不知道怎么中断回答
//                if (AiCommandManager.stopOutput.get()) {
//                    return;
//                }
                String content = message.getOutput().getChoices().get(0).getMessage().getContent();
                AiCommandManager.appendMessage(content);
                PrintUtil.print(content);
            }

            @Override
            public void onError(Exception err) {
                PrintUtil.printLnError("错误" + err.getMessage());
                semaphore.release();
            }

            @Override
            public void onComplete() {
                System.out.println("AI回答完毕");
                PrintUtil.println();
                semaphore.release();
            }
        });

        semaphore.acquire();
    }

    private static GenerationParam buildGenerationParam(List<Message> userMsgList) {
        // 打印配置信息
        AiConfig aiConfig = AiCommandManager.aiConfig;
        GenerationParam param = GenerationParam.builder()
                .model(aiConfig.getMode())
                .apiKey(getApiKey())
                .resultFormat(GenerationParam.ResultFormat.MESSAGE)
                .topP(aiConfig.getTopP())
                .incrementalOutput(true)
                .enableSearch(true)
                .build();

        if (aiConfig.getCurrentPrompt() != null) {
            PrintUtil.println("当前使用的Prompt：" + aiConfig.getCurrentPrompt().getName());
            PrintUtil.println();
            String content = userMsgList.get(userMsgList.size() - 1).getContent();
            param.setPrompt(aiConfig.getCurrentPrompt().getPrompt() + "\n" + content);
        } else {
            param.setMessages(userMsgList);

        }
        return param;
    }



    // 转换AiCommandManager的messageList
    public static List<Message> getMessageList() {
        List<String> messageList = AiCommandManager.getMessageList();
        List<Message> result = Arrays.asList(new Message[messageList.size()]);
        for (int i = 0; i < messageList.size(); i++) {
            Message userMsg = Message.builder().role(Role.USER.getValue()).content(messageList.get(i)).build();
            result.set(i, userMsg);
        }
        return result;
    }


}


